Graph RNNT: Grid- and Compose-Transducer. W-Transducer loss #6168

artbataev · 2023-03-10T16:49:35Z

What does this PR do?

Implement Grid- and Compose-Transducer along with W-Transducer loss according to the paper "Powerful and Extensible WFST Framework for Rnn-Transducer Losses" https://ieeexplore.ieee.org/document/10096679 (https://arxiv.org/abs/2303.10384)

Collection: [ASR]

Changelog

GraphTransducerLossBase abstract class with the interface for Graph-based loses
RNN-T implementation in GraphRnntLoss with tests
W-Transducer implementation in GraphWTransducerLoss with tests
add GraphRnntLoss + GraphWTransducerLoss to RNN-T loss resolver

Usage

# training some model with Graph-RNNT loss
python examples/asr/asr_transducer/speech_to_text_rnnt_bpe.py \
    <other params> \
    model.loss.loss_name=graph_rnnt \
    ++model.loss.graph_rnnt_kwargs.double_scores=true

# training some model with W-Transducer loss
python examples/asr/asr_transducer/speech_to_text_rnnt_bpe.py \
    <other params> \
    model.loss.loss_name=graph_w_transducer \
    ++model.loss.graph_w_transducer_kwargs.eps_weight=0

Before your PR is "Ready for review"

Pre checks:

Make sure you read and followed Contributor guidelines
Did you write any new necessary tests?
Did you add or update any necessary documentation?
Does the PR affect components that are optional to install? (Ex: Numba, Pynini, Apex etc)
- Reviewer: Does the PR have correct import guards for all optional libraries?

PR Type:

New Feature
Bugfix
Documentation

If you haven't finished some of the above items you can still open "Draft" PR.

Who can review?

Anyone in the community is free to review the PR once the checks have passed.
Contributor guidelines contains specific people who can review PRs to various areas.

Additional Information

Related to # (issue)

Signed-off-by: Vladimir Bataev <vbataev@nvidia.com>

nemo/collections/asr/losses/rnnt.py

github-actions · 2023-04-11T01:50:36Z

This PR is stale because it has been open for 14 days with no activity. Remove stale label or comment or update or this will be closed in 7 days.

github-actions · 2023-05-06T01:48:17Z

This PR is stale because it has been open for 14 days with no activity. Remove stale label or comment or update or this will be closed in 7 days.

Signed-off-by: Vladimir Bataev <vbataev@nvidia.com>

GNroy

Overall LGTM, please see comments.

nemo/collections/asr/parts/k2/graph_transducer.py

nemo/collections/asr/parts/k2/w_transducer.py

nemo/collections/asr/parts/k2/graph_transducer.py

nemo/collections/asr/parts/k2/w_transducer.py

Signed-off-by: Vladimir Bataev <vbataev@nvidia.com>

… into transucer_compose_grid_wildcard

Signed-off-by: Vladimir Bataev <vbataev@nvidia.com>

titu1994

Looks great. Minor comment about the new arg, make it explicit rather than a default. Easy to miss otherwise

titu1994 · 2023-05-26T18:31:34Z

nemo/collections/asr/losses/rnnt.py

@@ -109,6 +119,20 @@ class RNNTLossConfig:
        is_available=True,
        installation_msg="Pure Pytorch implementation of Multiblank RNN-T loss. Slow and for debugging purposes only.",
    ),
+    "graph_w_transducer": RNNTLossConfig(


Update the rest of the configs to explicitly show force fp32 is true

Fixed (I updated all the configs)

Signed-off-by: Vladimir Bataev <vbataev@nvidia.com>

) * add GraphTransducerLossBase abstract class with the interface for Graph-based loses * add RNN-T implementation in GraphRnntLoss with tests * add W-Transducer implementation in GraphWTransducerLoss with tests * add GraphRnntLoss + GraphWTransducerLoss to RNN-T loss resolver --------- Signed-off-by: Vladimir Bataev <vbataev@nvidia.com> Signed-off-by: hsiehjackson <c2hsieh@ucsd.edu>

artbataev added 2 commits March 10, 2023 20:46

Add graph transducer: grid and compose transducer

a265fb2

Signed-off-by: Vladimir Bataev <vbataev@nvidia.com>

Add w-transducer

10bd265

Signed-off-by: Vladimir Bataev <vbataev@nvidia.com>

github-actions bot added the ASR label Mar 10, 2023

Add Graph-RNNT to RNNT loss resolver

a1edc5d

Signed-off-by: Vladimir Bataev <vbataev@nvidia.com>

github-advanced-security bot found potential problems Mar 10, 2023

View reviewed changes

nemo/collections/asr/losses/rnnt.py Fixed Show fixed Hide fixed

GNroy mentioned this pull request Mar 13, 2023

[ASR] k2: Epsilon-Transducer and code imprevements #3962

Merged

8 tasks

artbataev added 2 commits March 22, 2023 19:58

Merge branch 'main' into transucer_compose_grid_wildcard

549ffc1

Merge branch 'main' into transucer_compose_grid_wildcard

5055d77

github-actions bot added the stale label Apr 11, 2023

GNroy removed the stale label Apr 11, 2023

Merge branch 'main' into transucer_compose_grid_wildcard

969b9f8

github-actions bot added the stale label May 6, 2023

artbataev removed the stale label May 7, 2023

artbataev added 6 commits May 7, 2023 21:15

Merge branch 'main' into transucer_compose_grid_wildcard

b86c4d2

Merge branch 'main' into transucer_compose_grid_wildcard

187053f

Merge branch 'main' into transucer_compose_grid_wildcard

6865d10

Temporary fix for k2 installation

240a566

Signed-off-by: Vladimir Bataev <vbataev@nvidia.com>

Improve graph-transducer: docstrings, refactoring

e024776

Signed-off-by: Vladimir Bataev <vbataev@nvidia.com>

Add graph-transducer losses to rnnt loss resolver

646e0b8

Signed-off-by: Vladimir Bataev <vbataev@nvidia.com>

github-actions bot added the core Changes to NeMo Core label May 19, 2023

artbataev added 8 commits May 19, 2023 19:42

Unify interface

9e32838

Signed-off-by: Vladimir Bataev <vbataev@nvidia.com>

Test Graph-RNNT

8088e33

Signed-off-by: Vladimir Bataev <vbataev@nvidia.com>

Fix device usage in tests

202d581

Signed-off-by: Vladimir Bataev <vbataev@nvidia.com>

Fix W-Transducer, add tests

efeeb7e

Signed-off-by: Vladimir Bataev <vbataev@nvidia.com>

Merge branch 'main' into transucer_compose_grid_wildcard

2c67b61

Switch to k2 original repo (fix for CUDA 12)

89b4c92

Signed-off-by: Vladimir Bataev <vbataev@nvidia.com>

Fix docstring

726e62f

Signed-off-by: Vladimir Bataev <vbataev@nvidia.com>

Fix docstring

a0758bf

Signed-off-by: Vladimir Bataev <vbataev@nvidia.com>

artbataev added 7 commits May 24, 2023 17:31

W-Transducer: test grid.

c58f8c8

Signed-off-by: Vladimir Bataev <vbataev@nvidia.com>

Merge branch 'main' into transucer_compose_grid_wildcard

41646dc

Improve docstrings.

2761b12

Signed-off-by: Vladimir Bataev <vbataev@nvidia.com>

Improve docstrings. Fix tests.

8bcbd54

Signed-off-by: Vladimir Bataev <vbataev@nvidia.com>

Improve docstrings

ccc014c

Signed-off-by: Vladimir Bataev <vbataev@nvidia.com>

Fix loss name in RNN-T resolver

276751a

Signed-off-by: Vladimir Bataev <vbataev@nvidia.com>

Merge branch 'main' into transucer_compose_grid_wildcard

f185bc2

artbataev requested a review from GNroy May 24, 2023 14:12

artbataev marked this pull request as ready for review May 24, 2023 14:12

artbataev added 4 commits May 24, 2023 20:56

Add quotes to type annotations

2330641

Signed-off-by: Vladimir Bataev <vbataev@nvidia.com>

scheme -> schema

6ed956e

Signed-off-by: Vladimir Bataev <vbataev@nvidia.com>

Merge branch 'main' into transucer_compose_grid_wildcard

f2f56aa

Merge branch 'main' into transucer_compose_grid_wildcard

3a327cd

GNroy previously approved these changes May 26, 2023

View reviewed changes

artbataev added 2 commits May 26, 2023 18:03

force_last -> force_final

7f43fa6

Signed-off-by: Vladimir Bataev <vbataev@nvidia.com>

Merge remote-tracking branch 'origin/transucer_compose_grid_wildcard'…

7a75f37

… into transucer_compose_grid_wildcard

artbataev dismissed GNroy’s stale review via 7a75f37 May 26, 2023 14:03

artbataev added 2 commits May 26, 2023 18:39

Add comments in forward method

d32e0e3

Signed-off-by: Vladimir Bataev <vbataev@nvidia.com>

W-Transducer: fix temporal scheme construction (remove redundant arc)

b0af4c2

Signed-off-by: Vladimir Bataev <vbataev@nvidia.com>

GNroy previously approved these changes May 26, 2023

View reviewed changes

artbataev requested a review from titu1994 May 26, 2023 15:57

titu1994 previously approved these changes May 26, 2023

View reviewed changes

RNNT_LOSS_RESOLVER: set force_float32 explicitly for all losses

607601e

Signed-off-by: Vladimir Bataev <vbataev@nvidia.com>

artbataev dismissed stale reviews from titu1994 and GNroy via 607601e May 26, 2023 18:56

GNroy approved these changes May 26, 2023

View reviewed changes

artbataev merged commit f1b5eae into main May 26, 2023

artbataev deleted the transucer_compose_grid_wildcard branch May 26, 2023 20:55

artbataev mentioned this pull request May 26, 2023

TDT model pull request #6536

Merged

8 tasks

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Graph RNNT: Grid- and Compose-Transducer. W-Transducer loss #6168

Graph RNNT: Grid- and Compose-Transducer. W-Transducer loss #6168

artbataev commented Mar 10, 2023 •

edited

Loading

github-actions bot commented Apr 11, 2023

github-actions bot commented May 6, 2023

GNroy left a comment

titu1994 left a comment

titu1994 May 26, 2023

artbataev May 26, 2023

Graph RNNT: Grid- and Compose-Transducer. W-Transducer loss #6168

Graph RNNT: Grid- and Compose-Transducer. W-Transducer loss #6168

Conversation

artbataev commented Mar 10, 2023 • edited Loading

What does this PR do?

Changelog

Usage

Before your PR is "Ready for review"

Who can review?

Additional Information

github-actions bot commented Apr 11, 2023

github-actions bot commented May 6, 2023

GNroy left a comment

Choose a reason for hiding this comment

titu1994 left a comment

Choose a reason for hiding this comment

titu1994 May 26, 2023

Choose a reason for hiding this comment

artbataev May 26, 2023

Choose a reason for hiding this comment

artbataev commented Mar 10, 2023 •

edited

Loading